import starlette.websockets
from server import login
from server.userobj import User
from fastapi import HTTPException
import fastapi
from typing import Literal
import aiohttp
import asyncio
import starlette
from urllib.parse import urlencode, urlunparse
# 绑定路由


async def permission_check_func(user: User, plug_name, level):
    if (await user.get_permission("//plugin/"+plug_name) < level):
        raise HTTPException(401)


async def forward_request(url: str, request: fastapi.Request):
    async with aiohttp.ClientSession() as session:
        headers = {key: value for key, value in request.headers.items()}
        headers['Accept-Encoding'] = 'identity'
        method = request.method
        body = await request.body()
        async with session.request(method, url, headers=headers, data=body) as resp:
            resp_body = await resp.read()
            return resp.status, resp.headers, resp_body


class SetConfigExit(Exception):
    def __init__(self, message="", *args):
        super().__init__(message, *args)


class Route():
    def __init__(self, plugid: str):
        # 设置信息列表
        self.plugid = plugid
        self.apiroot = "/api/plugin/"+self.plugid+"/api"
        self.webroot = "/api/plugin/"+self.plugid+"/iframe"
        self.static = "/api/plugin/"+self.plugid+"/static"
        self.routes_get = {}
        self.routes_post = {}
        self.routes_delete = {}
        self.routes_ws = {}
        self.routes_proxy = {}

    def get(self, route: str, level: Literal[0, 1, 2] = 1):
        # 注册get
        def wraproot(func):
            # 修饰器修补参数添加校验
            warpper = login.check_cookie(
                func, permission_check_func, self.plugid, level=level)
            self.routes_get[route] = warpper
            return warpper
        return wraproot

    def post(self, route: str, level: Literal[0, 1, 2] = 2):
        def wraproot(func):
            # 修饰器修补参数添加校验
            warpper = login.check_cookie(
                func, permission_check_func, self.plugid, level=level)
            self.routes_post[route] = warpper
            return warpper
        return wraproot

    def delete(self, route: str, level: Literal[0, 1, 2] = 2):
        def wraproot(func):
            # 修饰器修补参数添加校验
            warpper = login.check_cookie(
                func, permission_check_func, self.plugid, level=level)
            self.routes_delete[route] = warpper
            return warpper
        return wraproot

    def ws(self, route, level=2):
        def wraproot(func):
            # 修饰器修补参数添加校验
            warpper = login.check_cookie(
                func, permission_check_func, self.plugid, level=level)
            self.routes_ws[route] = warpper
            return warpper
        return wraproot

    def proxy(self, route: str, level: Literal[0, 1, 2] = 2):
        def wraproot(func):
            # 修饰器修补参数添加校验
            async def response_func(user: User, anypath: str, request: fastapi.Request):
                async def geturl(target: str):
                    target_url = target.rstrip("/")+"/" + anypath
                    if (request.url.query != ""):
                        target_url += "?" + str(request.url.query)
                    status, headers, body = await forward_request(target_url, request)
                    return (fastapi.Response(content=(body), status_code=status, headers=headers))
                try:
                    return await func(user, geturl)
                except aiohttp.ClientConnectorError:
                    return None
            warpper = login.check_cookie(
                response_func, permission_check_func, self.plugid, level=level)
            self.routes_proxy[route+"/{anypath:path}"] = warpper

            async def ws_response_func(user: User, anypath: str, websocket: fastapi.WebSocket):
                target_url = ""

                async def geturl(target: str):
                    nonlocal target_url
                    target_url = target.rstrip("/")+"/" + anypath
                    if (websocket.url.query != ""):
                        target_url += "?" + str(websocket.url.query)
                    raise SetConfigExit
                try:
                    await func(user, geturl)
                except SetConfigExit:
                    pass
                await websocket.accept()

                # 将客户端的请求头转换为 aiohttp 所需的格式
                extra_headers = {key: value for key,
                                 value in websocket.headers.items() if key.lower().startswith("x-") or key.lower() == "user-agent" or key.lower() == "cookie"}
                try:
                    async with aiohttp.ClientSession() as session:
                        async with session.ws_connect(target_url, headers=extra_headers) as target_ws:
                            # 双向转发消息
                            async def forward_messages_consumer(source: fastapi.WebSocket, target: aiohttp.ClientWebSocketResponse):
                                try:
                                    while True:
                                        msg = await source.receive_bytes()
                                        await target.send_bytes(msg)
                                except starlette.websockets.WebSocketDisconnect:
                                    pass

                            async def forward_messages_producer(source: aiohttp.ClientWebSocketResponse, target: fastapi.WebSocket):
                                async for msg in source:
                                    if msg.type == aiohttp.WSMsgType.TEXT:
                                        await target.send_text(msg.data)
                                    elif msg.type == aiohttp.WSMsgType.BINARY:
                                        await target.send_bytes(msg.data)
                                    elif msg.type == aiohttp.WSMsgType.CLOSE:
                                        break
                                    elif msg.type == aiohttp.WSMsgType.ERROR:
                                        break

                            # 从客户端到目标服务器
                            consumer_task = asyncio.create_task(
                                forward_messages_consumer(websocket, target_ws))

                            # 从目标服务器到客户端
                            producer_task = asyncio.create_task(
                                forward_messages_producer(target_ws, websocket))

                            done, pending = await asyncio.wait(
                                [consumer_task, producer_task],
                                return_when=asyncio.FIRST_COMPLETED,
                            )

                            for task in pending:
                                task.cancel()
                except aiohttp.ClientConnectorError:
                    pass
            warpper2 = login.check_cookie(
                ws_response_func, permission_check_func, self.plugid, level=level)
            self.routes_ws[route+"/{anypath:path}"] = warpper2
            return warpper
        return wraproot
